﻿#include "LennardJonesSystem.h"
#include "Application.h"
#include "Mesh.h"
#include "Transient.h"

#include <vtkm/Pair.h>
#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/Algorithm.h>
#include <vtkm/cont/ArrayHandleGroupVecVariable.h>
#include <vtkm/cont/ConvertNumComponentsToOffsets.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

struct LennardJonesSystem::InitPositionWorklet : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut);
  using ExecutionSignature = void(_2, _3);

  template<typename PType>
  VTKM_EXEC void operator()(const PType& pts, vtkm::Vec3f& centroid) const
  {
    if (pts.GetNumberOfComponents() == 4)
    {
      centroid = (pts[0] + pts[1] + pts[2] + pts[3]) / 4;
    }
    else if (pts.GetNumberOfComponents() == 8)
    {
      centroid = (pts[0] + pts[1] + pts[2] + pts[3] + pts[4] + pts[5] + pts[6] + pts[7]) / 8;
    }
    else
    {
      RaiseError("unfinished");
    }
  }
};

struct LennardJonesSystem::InitialCondtionWorklet : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn pts, FieldOut mass, FieldOut veclocity);
  using ExecutionSignature = void(_1, _2, _3);

  template<typename CoordType>
  VTKM_EXEC void operator()(const CoordType& pts, Real& mass, Vec3f& velocity) const
  {
    mass = 1.;
    velocity = { 0.01, 0.01, 0 };
  }
};

struct LennardJonesSystem::ComputeForceWorklet : vtkm::worklet::WorkletMapField
{
  ComputeForceWorklet(const vtkm::Vec3f& dxdydz)
    : _dxdydz(dxdydz)
  {
  }

  using ControlSignature = void(FieldIn current_pts,
                                FieldIn neighbor_index,
                                FieldOut force,
                                WholeArrayIn whole_pts);
  using ExecutionSignature = void(_1, _2, _3, _4);

  template<typename CoordType, typename ArrayCoordType, typename WholePtsType>
  VTKM_EXEC void operator()(const CoordType& current_pts,
                            const ArrayCoordType& neighbor_index,
                            Vec3f& force,
                            WholePtsType& whole_pts) const
  {
    force = { 0, 0, 0 };
    const auto& p_i = current_pts;
    Real sigma_2 = _sigma * _sigma;
    Real sigma_6 = sigma_2 * sigma_2 * sigma_2;
    Real sigma_12 = sigma_6 * sigma_6;
    for (size_t j = 0; j < neighbor_index.GetNumberOfComponents(); j++)
    {
      auto offset = neighbor_index[j].Get().second;
      const auto& p_j = whole_pts.Get(neighbor_index[j].Get().first) +
        vtkm::Vec3f(offset[0] * _dxdydz[0], offset[1] * _dxdydz[1], offset[2] * _dxdydz[2]);
      auto r_ij = p_i - p_j;
      Real dis_2 = vtkm::MagnitudeSquared(r_ij);
      if (dis_2 <= 0.001)
        continue;
      Real dis_6 = dis_2 * dis_2 * dis_2;
      Real dis_12 = dis_6 * dis_6;
      force -= (2 * sigma_12 / dis_12 - sigma_6 / dis_6) / dis_2 * r_ij;
    }
    force = force * 24 * _eps;
  }

  Real _eps = 1;
  Real _sigma = 1;
  vtkm::Vec3f _dxdydz;
};

struct LennardJonesSystem::ComputeForceWithPeriodicBC : vtkm::worklet::WorkletMapField
{
  ComputeForceWithPeriodicBC(Real rs)
    : _rs2(rs*rs)
  {
  }

  using ControlSignature = void(FieldIn current_pts,
                                FieldOut force,
                                WholeArrayIn whole_pts,
                                ExecObject locator);
  using ExecutionSignature = void(_1, _2, _3, _4);

  template<typename CoordType, typename WholePtsType>
  VTKM_EXEC void operator()(const CoordType& current_pts,
                            Vec3f& force,
                            WholePtsType& whole_pts,
                            const SCIAI::exec::PointLocatorSparseGrid& locator) const
  {
    force = { 0, 0, 0 };
    const auto& p_i = current_pts;
    auto current_ijk = locator.InWhichCell(p_i);
    //std::cout << "current_pts: " << current_pts  << " current_ijk: " << current_ijk << std::endl;
    int status = 0;
    for (Id i = -1; i <= 1; i++)
    {
      for (Id j = -1; j <= 1; j++)
      {
        for (Id k = -1; k <= 1; k++)
        {
          auto neighbor_ijk = current_ijk + Id3{ i, j, k };
          auto ijk = locator.PeriodicIndexOffset(neighbor_ijk);
          auto coord_offset = locator.PeriodicCoordOffset(neighbor_ijk);
          auto num_pts = locator.NumberPtsInCell(ijk);
          //std::cout << "neighbor_ijk: " << neighbor_ijk << std::endl;
          //std::cout << "ijk: " << ijk << std::endl;
          //std::cout << "coord_offset: " << coord_offset << std::endl;
          //std::cout << "num_pts: " << num_pts << std::endl;
          for (Id p = 0; p < num_pts; p++)
          {
            auto pts_id = locator.PtsInCell(ijk, p);
            auto p_j = whole_pts.Get(pts_id) - coord_offset;

            force += ComputeEnergyForce(p_i, p_j);
            //auto r_ij = p_i - p_j;
            //Real dis_2 = vtkm::MagnitudeSquared(r_ij);
            //std::cout << "p_i: " << p_i << " p_j: " << p_j << " dis2:" << dis_2 << std::endl;

            /*       if ( dis_2 < 1.5*1.5)
            {
              status++;
            }*/
          }
        }
      }
    }

    force = force * 24 * _eps;
    //std::cout << force << std::endl;

    //std::cout << "current_pts: " << current_pts << " 邻近点个数：" << status << std::endl;
    //std::exit(0);
  }

  VTKM_EXEC Vec3f ComputeEnergyForce(const Vec3f& p_i, const Vec3f& p_j) const
  {
    Vec3f force{ 0, 0, 0 };
    auto r_ij = p_j - p_i;
    Real dis_2 = vtkm::MagnitudeSquared(r_ij);
    //std::cout << "p_i: " << p_i << "p_j: " << p_j << " dis2:" <<  dis_2 << std::endl;
    if (dis_2 > 0.01 && dis_2 < _rs2)
    {
      Real dis_6 = dis_2 * dis_2 * dis_2;
      Real dis_12 = dis_6 * dis_6;
      force = -(2 * _sigma_12 / dis_12 - _sigma_6 / dis_6) / dis_2 * r_ij;
      //std::cout << "p_i: " << p_i << "p_j: " << p_j << " dis2:" <<  dis_2 << std::endl;
      //std::cout << force << std::endl;
    }
    return force;
  }

  Real _rs2; //截断距离
  Real _eps = 1;
  Real _sigma = 1;
  Real _sigma_2 = _sigma * _sigma;
  Real _sigma_6 = _sigma_2 * _sigma_2 * _sigma_2;
  Real _sigma_12 = _sigma_6 * _sigma_6;

  //vtkm::Vec3f _dxdydz;
};

struct LennardJonesSystem::UpdateVelocityWorklet : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldInOut velocity, FieldIn force, FieldIn mass);
  using ExecutionSignature = void(_1, _2, _3);

  UpdateVelocityWorklet(Real dt)
    : _dt(dt)
  {
  }

  VTKM_EXEC void operator()(Vec3f& velocity, const Vec3f& force, const Real& mass) const
  {
    velocity += 0.5 * force / mass * _dt;
  }

  Real _dt;
};

struct LennardJonesSystem::UpdatePositionWorklet : vtkm::worklet::WorkletMapField
{
  UpdatePositionWorklet(Real dt)
    : _dt(dt)
  {
  }

  using ControlSignature = void(FieldInOut position, FieldIn velocity, ExecObject locator);
  using ExecutionSignature = void(_1, _2, _3);

  template<typename CoordType>
  VTKM_EXEC void operator()(CoordType& position,
                            const Vec3f& velocity,
                            SCIAI::exec::PointLocatorSparseGrid locator) const
  {
    position += velocity * _dt;
    locator.UpdateOverRangePoint(position);
  }
  Real _dt;
};

class LennardJonesSystem::ComputeNeighborPtsOffsetWorklet : public vtkm::worklet::WorkletMapField
{
public:
  typedef void ControlSignature(FieldIn qcIn, FieldIn disIn, ExecObject locator, FieldOut nnIdOut);

  typedef void ExecutionSignature(_1, _2, _3, _4);

  VTKM_CONT
  ComputeNeighborPtsOffsetWorklet() {}

  template<typename CoordiVecType, typename DisType, typename Locator>
  VTKM_EXEC void operator()(const CoordiVecType& qc,
                            const DisType& dis,
                            const Locator& locator,
                            vtkm::Id& pointNumber) const
  {
    locator.FindNearestNeighbor(qc, dis, pointNumber);
  }
};


class LennardJonesSystem::ComputeNeighborPtsIndexWorklet : public vtkm::worklet::WorkletMapField
{
public:
  typedef void ControlSignature(FieldIn qcIn, FieldIn disIn, ExecObject locator, FieldInOut ids);

  typedef void ExecutionSignature(_1, _2, _3, _4);

  VTKM_CONT
  ComputeNeighborPtsIndexWorklet() {}

  template<typename CoordiVecType, typename DisType, typename Locator, typename ConnectivityVecType>
  VTKM_EXEC void operator()(const CoordiVecType& qc,
                            const DisType& dis,
                            const Locator& locator,
                            ConnectivityVecType& ids) const
  {
    locator.FindPointIds(qc, dis, ids);
  }
};


RegisterObject(LennardJonesSystem);

LennardJonesSystem::LennardJonesSystem(const Configuration& cfg)
  : System(cfg)
  , _dis(cfg.Get<Real>("dis"))
  , _bin_number(cfg.GetVectorValue<int>("bin_number"))

{
  auto locator_x_range = cfg.GetVectorValue<Real>("locator_x_range");
  auto locator_y_range = cfg.GetVectorValue<Real>("locator_y_range");
  auto locator_z_range = cfg.GetVectorValue<Real>("locator_z_range");
  _locator_range = vtkm::Vec<vtkm::Range, 3>(vtkm::Range(locator_x_range[0], locator_x_range[1]),
                                             vtkm::Range(locator_y_range[0], locator_y_range[1]),
                                             vtkm::Range(locator_z_range[0], locator_z_range[1]));
  _rs = _dis;
}

void LennardJonesSystem::Init()
{
  System::Init();
  _data_set.SetCellSet(_mesh.GetCellSet());
  _data_set.AddCoordinateSystem(_mesh.GetCoordinate());
  //this->AddCellVariable<Vec3f>("velocity");
  //this->AddCellVariable<Real>("mass");
  //this->AddCellVariable<Vec3f>("force");

  // 设置 velocity和mass的初始值
  this->InitialCondition();

  //init locator
  InitPointLocator();

  //init dis
  auto point_number = this->_data_set.GetNumberOfPoints();
  _dis_array.AllocateAndFill(point_number, _dis);
}
void LennardJonesSystem::InitialCondition()
{
  vtkm::cont::Invoker{}(LennardJonesSystem::InitPositionWorklet{},
                        _data_set.GetCellSet(),
                        _data_set.GetCoordinateSystem(),
                        _position);

  //auto mass = this->GetVariableArray<Real>("mass");
  //auto velocity = this->GetVariableArray<Vec3f>("velocity");
  vtkm::cont::Invoker{}(LennardJonesSystem::InitialCondtionWorklet{}, _position, _mass, _velocity);
}

void LennardJonesSystem::ComputeForce()
{
  //auto force = this->GetVariableArray<Vec3f>("force");
  //auto dxdydz = vtkm::Vec3f(_locator_range[0].Max - _locator_range[0].Min,
  //                          _locator_range[1].Max - _locator_range[1].Min,
  //                          _locator_range[2].Max - _locator_range[2].Min);
  //vtkm::cont::Invoker{}(LennardJonesSystem::ComputeForceWorklet{ dxdydz },
  //                      _position,
  //                      _point_ids_array,
  //                      force,
  //                      _position);

  vtkm::cont::Invoker{}(
    LennardJonesSystem::ComputeForceWithPeriodicBC{ _rs }, _position, _force, _position, _locator);
}

void LennardJonesSystem::UpdateVelocity()
{
  try
  {
    // TODO: mass, velocity, force的ArrayHandle作为变量？
    //auto mass = this->GetVariableArray<Real>("mass");
    //auto velocity = this->GetVariableArray<Vec3f>("velocity");
    //auto force = this->GetVariableArray<Vec3f>("force");

    _invoker(LennardJonesSystem::UpdateVelocityWorklet{ _dt }, _velocity, _force, _mass);
  }
  catch (const std::exception& e)
  {
    std::cout << e.what() << std::endl;
  }
}

void LennardJonesSystem::UpdatePosition()
{
  //auto velocity = this->GetVariableArray<Vec3f>("velocity");
  _invoker(LennardJonesSystem::UpdatePositionWorklet{ _dt }, _position, _velocity, _locator);
}


void LennardJonesSystem::Evolve()
{
  _timer.Start();

  // Get Time Step: dt
  PreSolve();

  //stage0:
  //ComputeNeighborPtsOffset();
  //ComputeNeighborPtsIndex();

  // stage1:
  ComputeForce();
  UpdateVelocity();

  // stage2:
  UpdatePosition();

  // stage3:
  ComputeForce();
  UpdateVelocity();

  PostSolve();
}

void LennardJonesSystem::PreSolve()
{
  _locator.Update();
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  _dt = transient.Dt();
}



void LennardJonesSystem::TimeIntegration() {}

void LennardJonesSystem::Solve() {}

void LennardJonesSystem::PostSolve()
{
  //auto force = this->GetVariableArray<Vec3f>("force");
  auto result =
    vtkm::cont::Algorithm::Reduce(_force, vtkm::TypeTraits<Vec3f>::ZeroInitialization());
  Info("LennardJonesSystem::Evolve运行完成，", "计算时间: ", _timer.GetElapsedTime(), " s");
  Info("LennardJonesSystem::force ", result);
  _data_set.AddCellField("force", _force);
  _data_set.AddCellField("velocity", _velocity);
}

void LennardJonesSystem::InitPointLocator()
{
  //set coord
  //auto coord = this->_data_set.GetCoordinateSystem();
  _locator.SetCoordinates(CoordinateSystem("position", _position));

  //set range
  _locator.SetRange(_locator_range);

  //set number of bins
  _locator.SetNumberOfBins(vtkm::Id3{ _bin_number[0], _bin_number[1], _bin_number[2] });

  //update
  _locator.Update();
}

void LennardJonesSystem::ComputeNeighborPtsOffset()
{
  auto coord = _data_set.GetCoordinateSystem();
  ArrayHandle<Id> point_number_array;
  _invoker(ComputeNeighborPtsOffsetWorklet{}, coord, _dis_array, _locator, point_number_array);
  auto sum_points_number = vtkm::cont::Algorithm::Reduce(
    point_number_array, vtkm::TypeTraits<vtkm::Id>::ZeroInitialization());

  //auto number_portal = point_number_array.ReadPortal();
  //for (int i = 0; i < number_portal.GetNumberOfValues(); ++i)
  //{
  //  std::cout << number_portal.Get(i) << std::endl;
  //}
  //std::cout << "sum_points_number: " << sum_points_number << std::endl;

  //group point ids
  vtkm::cont::ArrayHandle<vtkm::Pair<vtkm::Id, vtkm::Vec<vtkm::Int8, 3>>> neighbor_pts_index;
  neighbor_pts_index.Allocate(sum_points_number);
  auto neighbor_pts_offset = vtkm::cont::ConvertNumComponentsToOffsets(point_number_array);
  _point_ids_array =
    vtkm::cont::make_ArrayHandleGroupVecVariable(neighbor_pts_index, neighbor_pts_offset);
}

void LennardJonesSystem::ComputeNeighborPtsIndex()
{
  //fill point ids
  auto coord = _data_set.GetCoordinateSystem();
  _invoker(ComputeNeighborPtsIndexWorklet{}, coord, _dis_array, _locator, _point_ids_array);

  //auto ids_portal = _point_ids_array.ReadPortal();
  //for (int i = 0; i < ids_portal.GetNumberOfValues(); ++i)
  //{
  //  auto ids = ids_portal.Get(i);
  //  for (auto index = 0; index < ids.GetNumberOfComponents(); ++index)
  //  {
  //    std::cout << "first: " << ids[index].Get().first
  //              << "   second: " << +ids[index].Get().second[0] << +ids[index].Get().second[1]
  //              << +ids[index].Get().second[2] << std::endl;
  //  }
  //}
}
